"""
SQLAlchemy 数据库连接器
"""
from typing import Dict, Any, Union
from sqlalchemy import create_engine
from sqlalchemy.orm import sessionmaker
from sqlalchemy.engine import Engine
from .base_connector import BaseConnector


class SQLAlchemyConnector(BaseConnector):
    """
    SQLAlchemy 数据库连接器
    """
    def __init__(self):
        self.engine = None
        self.session_factory = None

    def connect(self, connection_params: Dict[str, Any]) -> Union[Engine, sessionmaker]:
        """
        连接数据库
        
        Args:
            connection_params: 连接参数
        Returns:
            SQLAlchemy Engine 或 Session 工厂
        """
        db_url = f"{connection_params.get('dialect')}://{connection_params.get('username')}:{connection_params.get('password')}@{connection_params.get('host')}:{connection_params.get('port')}/{connection_params.get('database')}"
        self.engine = create_engine(db_url)
        self.session_factory = sessionmaker(bind=self.engine)
        return self.engine

    def execute_query(self, query: str) -> Dict[str, Any]:
        """
        执行 SQL 查询
        
        Args:
            query: SQL 查询语句
        """
        with self.engine.connect() as connection:
            result = connection.execute(query)
            return [dict(row) for row in result]

    def close(self) -> None:
        """
        关闭数据库连接
        """
        if self.engine:
            self.engine.dispose()